import numpy as np
from autom.prepro import translate,get_part
# 搜索填充方向
def search_overlap1(M):
    m,n = 100,100
    lower_bound = 10
    sim = np.zeros((m-lower_bound,n-lower_bound))
    for py in range(lower_bound,m):
        for px in range(lower_bound,n):
            S1 = M[py:,px:]
            S2 = M[:-py,:-px]
            sim[py-lower_bound,px-lower_bound] = np.mean((S1-S2)**2)
    py,px = np.where(sim==np.min(sim))
    return py[0]+lower_bound,px[0]+lower_bound
def search_overlap2(M):
    m,n = 100,100
    lower_bound = 10
    sim = np.zeros((m-lower_bound,n-lower_bound))
    for py in range(-m,-lower_bound):
        for px in range(lower_bound,n):
            S1 = M[-py:,:M.shape[1]-px]
            S2 = M[:M.shape[0]+py,px:]
            sim[py+m,px-lower_bound] = np.mean((S1-S2)**2)
    py,px = np.where(sim==np.min(sim))
    return py[0]-m,px[0]+lower_bound
def search_overlap(M):
    py1,px1 = search_overlap1(M)
    py2,px2 = search_overlap2(M)
    return py1,px1,py2,px2

# 填充算法
def impute(extended,region,px1,py1,px2,py2,t1,t2):
    height,width = extended.shape
    region_new = translate_region(region,px1,py1,px2,py2,t1,t2)
    y_low,y_up,x_low,x_up = region_new[1],region_new[3],region_new[0],region_new[2]
    y_low = np.clip(y_low,0,height)
    y_up = np.clip(y_up,0,height)
    x_low = np.clip(x_low,0,width)
    x_up = np.clip(x_up,0,width)
    y_low,y_up,x_low,x_up = int(y_low),int(y_up),int(x_low),int(x_up)
    region_back = translate_region([x_low,y_low,x_up,y_up],px1,py1,px2,py2,-t1,-t2)
    extended[y_low:y_up,x_low:x_up] = extended[region_back[1]:region_back[3],region_back[0]:region_back[2]]
    return extended

def translate_region(region,px1,py1,px2,py2,t1,t2):
    # region + ti*(px1,py1,px2,py2)
    region_new = region.copy()
    region_new[0] += t1*px1+t2*px2
    region_new[2] += t1*px1+t2*px2
    region_new[1] += t1*py1+t2*py2
    region_new[3] += t1*py1+t2*py2
    return region_new

def extend(M,px1,py1,px2,py2,region,height,width):
    extended = np.zeros((height,width))
    extended[region[1]:region[3],region[0]:region[2]] = M
    max_t1 = np.minimum(np.abs(height//py1),np.abs(width//px1))
    max_t2 = np.minimum(np.abs(height//py2),np.abs(width//px2))
    for t1 in range(-max_t1,max_t1):
        for t2 in range(-max_t2,max_t2):
            extended = impute(extended,region,px1,py1,px2,py2,t1,t2)
    return extended

def ex_impute(img,region1,region2,return_cor=False,progress=None):
    height,width = img.shape
    #region1 = translate(img.shape[0],region1)
    #region2 = translate(img.shape[0],region2)
    img1 = get_part(img,region1)
    img2 = get_part(img,region2)
    if type(progress) != type(None):
        progress.value = 25
    py11,px11,py12,px12 = search_overlap(img1)
    if type(progress) != type(None):
        progress.value = 50
    extended1 = extend(img1,px11,py11,px12,py12,region1,height=height,width=width)
    if type(progress) != type(None):
        progress.value = 65
    py21,px21,py22,px22 = search_overlap(img2)
    if type(progress) != type(None):
        progress.value = 75
    extended2 = extend(img2,px21,py21,px22,py22,region2,height=height,width=width)
    if type(progress) != type(None):
        progress.value = 85
    if return_cor:
        return extended1,extended2,(py11,px11,py12,px12),(py21,px21,py22,px22)
    else:
        return extended1,extended2



